package cn.amossun.starter.mybatis.data.security.core;

import cn.hutool.core.collection.CollectionUtil;
import cn.hutool.core.util.StrUtil;
import cn.amossun.starter.mybatis.data.security.enums.SqlDmlTypeEnum;
import cn.amossun.starter.mybatis.data.security.rewrite.condition.DecryptEncryptCondition;
import cn.amossun.starter.mybatis.data.security.rewrite.condition.DmlDecryptEncryptCondition;
import cn.amossun.starter.mybatis.data.security.rule.MybatisEncryptRule;
import cn.amossun.starter.mybatis.data.security.visitor.TablesNamesFinderPlus;
import lombok.Data;
import lombok.extern.slf4j.Slf4j;
import net.sf.jsqlparser.JSQLParserException;
import net.sf.jsqlparser.expression.Expression;
import net.sf.jsqlparser.expression.Function;
import net.sf.jsqlparser.expression.JdbcParameter;
import net.sf.jsqlparser.expression.operators.relational.EqualsTo;
import net.sf.jsqlparser.parser.CCJSqlParserUtil;
import net.sf.jsqlparser.schema.Column;
import net.sf.jsqlparser.statement.Statement;
import net.sf.jsqlparser.statement.delete.Delete;
import net.sf.jsqlparser.statement.insert.Insert;
import net.sf.jsqlparser.statement.select.PlainSelect;
import net.sf.jsqlparser.statement.select.Select;
import net.sf.jsqlparser.statement.select.SelectBody;
import net.sf.jsqlparser.statement.select.SetOperationList;
import net.sf.jsqlparser.statement.update.Update;
import net.sf.jsqlparser.util.deparser.ExpressionDeParser;

import java.util.*;

@Slf4j
@Data
public class CCJSQLStatementContext extends AbstractStatementContext {

    private final MybatisEncryptRule encryptRule;

    public CCJSQLStatementContext(String sql, MybatisEncryptRule encryptRule) throws JSQLParserException {
        setSql(sql);

        Statement statement = CCJSqlParserUtil.parse(sql);
        setStatement(statement);


        TablesNamesFinderPlus tablesNamesFinderPlus = new TablesNamesFinderPlus();
        setTableList(tablesNamesFinderPlus.getTableList(statement));
        this.encryptRule = encryptRule;
    }

    public DmlDecryptEncryptCondition findEncryptConditions() throws Exception {

        Statement statements = getStatement();
        if (statements instanceof Insert) {
            Insert insert = (Insert) statements;
            return DmlDecryptEncryptCondition.builder().encryptConditions(encryptResult(parseColumn(insert.getColumns()))).sqlDmlTypeEnum(SqlDmlTypeEnum.INSERT).build();
        }

        if (statements instanceof Delete) {
            Delete delete = (Delete) statements;
            return DmlDecryptEncryptCondition.builder().encryptConditions(encryptResult(parseDelete(delete))).sqlDmlTypeEnum(SqlDmlTypeEnum.DELETE).build();
        }

        if (statements instanceof Update) {
            Update update = (Update) statements;
            return DmlDecryptEncryptCondition.builder().encryptConditions(encryptResult(parseUpdate(update))).sqlDmlTypeEnum(SqlDmlTypeEnum.UPDATE).build();
        }

        if (statements instanceof Select) {
            Select select = (Select) statements;
            return DmlDecryptEncryptCondition.builder().encryptConditions(encryptResult(parseSelect(select))).sqlDmlTypeEnum(SqlDmlTypeEnum.SELECT).build();
        }

        return null;
    }

    protected List<DecryptEncryptCondition> encryptResult(List<CCJSQLCryptExpression> ccjSqlCryptExpressionList) throws Exception {
        if (CollectionUtil.isEmpty(ccjSqlCryptExpressionList)){
            return Collections.emptyList();
        }
        List<DecryptEncryptCondition> result = new ArrayList<>();
        for (CCJSQLCryptExpression ccjSqlCryptExpression : ccjSqlCryptExpressionList){
            result.addAll(buildEncryptConditions(ccjSqlCryptExpression));
        }
        return result;
    }

    private List<CCJSQLCryptExpression> parseDelete(Delete delete) {
        return whereExpression(delete.getWhere());
    }

    private List<CCJSQLCryptExpression> parseUpdate(Update update) {
        List<CCJSQLCryptExpression> columnList = parseColumn(update.getColumns());
        List<CCJSQLCryptExpression> whereList = whereExpression(update.getWhere());

        List<CCJSQLCryptExpression> result = new ArrayList<>(columnList.size() + whereList.size());
        if(CollectionUtil.isNotEmpty(columnList)) {
            result.addAll(columnList);
        }
        if(CollectionUtil.isNotEmpty(whereList)) {
            result.addAll(whereList);
        }
        return result;
    }

    private List<CCJSQLCryptExpression> parseColumn(List<Column> columns) {
        List<CCJSQLCryptExpression> ccjSqlCryptExpressionList = new ArrayList<>(columns.size());
        for (int i = 0; i < columns.size(); i++) {
            Column column = columns.get(i);
            CCJSQLCryptExpression ccjSqlCryptExpression = new CCJSQLCryptExpression();
            ccjSqlCryptExpression.setAlias(TablesNamesFinderPlus.getAlias(column.getTable()));
            ccjSqlCryptExpression.setColumnName(column.getColumnName());
            ccjSqlCryptExpression.setIndex(i + 1);
            ccjSqlCryptExpressionList.add(ccjSqlCryptExpression);
        }
        return ccjSqlCryptExpressionList;
    }

    public List<CCJSQLCryptExpression> parseSelect(Select select) {
        return parseSelectBody(select.getSelectBody());
    }

    protected List<CCJSQLCryptExpression> parseSelectBody(SelectBody selectBody) {
        if(selectBody instanceof  PlainSelect) {
            PlainSelect plain = (PlainSelect) selectBody;
            return whereExpression(plain.getWhere());
        }
        if(selectBody instanceof SetOperationList) {
            SetOperationList setOperationList = (SetOperationList) selectBody;
            if(CollectionUtil.isNotEmpty(setOperationList.getSelects())) {
                List<CCJSQLCryptExpression> ccjSqlCryptExpressionList = new ArrayList<>(16);
                setOperationList.getSelects().forEach(childSelectBody->{
                    ccjSqlCryptExpressionList.addAll(parseSelectBody(childSelectBody));
                });
                return ccjSqlCryptExpressionList;
            }
        }
        return null;
    }

    protected List<CCJSQLCryptExpression> whereExpression(Expression where) {
        if(where == null){
            return null;
        }
        StringBuilder buffer = new StringBuilder();
        List<CCJSQLCryptExpression> ccjSqlCryptExpressionList = new ArrayList<>(10);
        where.accept(new ExpressionDeParser(null, buffer) {
            @Override
            public void visit(EqualsTo equalsTo) {
                Expression leftExpression = equalsTo.getLeftExpression();
                Expression rightExpression = equalsTo.getRightExpression();

                if (leftExpression instanceof Column && rightExpression instanceof JdbcParameter){
                    Column column = (Column) equalsTo.getLeftExpression();
                    JdbcParameter jdbcParameter = (JdbcParameter) rightExpression;
                    CCJSQLCryptExpression param = new CCJSQLCryptExpression();
                    param.setIndex(jdbcParameter.getIndex());
                    param.setAlias(TablesNamesFinderPlus.getAlias(column.getTable()));
                    param.setColumnName(column.getColumnName());
                    ccjSqlCryptExpressionList.add(param);
                }

                if(leftExpression instanceof Function && rightExpression instanceof JdbcParameter) {
                    Function function = (Function) equalsTo.getLeftExpression();
                    JdbcParameter jdbcParameter = (JdbcParameter) rightExpression;

                }
            }
        });
        return ccjSqlCryptExpressionList;
    }

    private Collection<DecryptEncryptCondition> buildEncryptConditions(CCJSQLCryptExpression ccjSqlCryptExpression) throws Exception {
        Collection<DecryptEncryptCondition> result = new ArrayList<>();
        buildEncryptCondition(ccjSqlCryptExpression).ifPresent(result::add);
        return result;
    }

    private Optional<DecryptEncryptCondition> buildEncryptCondition(CCJSQLCryptExpression ccjSqlCryptExpression) throws Exception {
        Optional<String> tableName = findTableName(ccjSqlCryptExpression);
        final DecryptEncryptCondition condition = new DecryptEncryptCondition(tableName.isPresent() ? tableName.get() : null, ccjSqlCryptExpression.getColumnName(), ccjSqlCryptExpression.getIndex());
        return tableName.isPresent() && containsColumn(tableName.get(), ccjSqlCryptExpression.getColumnName()) ? Optional.of(condition) : Optional.empty();
    }

    public Optional<String> findTableName(CCJSQLCryptExpression ccjSqlCryptExpression) throws Exception {
        if (getTableList().size() == 1){
            return Optional.of(CollectionUtil.getFirst(getTableList()).getTableName());
        } else if (StrUtil.isNotEmpty(ccjSqlCryptExpression.getAlias())){
            return Optional.of(findTableNameFromSQL(ccjSqlCryptExpression.getAlias()));
        }

        return findTableNameFromMetaData(ccjSqlCryptExpression.getColumnName());
    }

    private Optional<String> findTableNameFromMetaData(final String columnName) throws Exception {
        for (CCJSQLTable each : getTableList()) {
            if (containsColumn(each.getTableName(), columnName)) {
                return Optional.of(each.getTableName());
            }
        }
        return Optional.empty();
    }

    private boolean containsColumn(String tableName, String columnName) throws Exception {
        return encryptRule.containsColumn(tableName, columnName);
    }

    private String findTableNameFromSQL(final String tableNameOrAlias) {
        for (CCJSQLTable ccjsqlTable : getTableList()) {
            if (tableNameOrAlias.equalsIgnoreCase(ccjsqlTable.getTableName()) || tableNameOrAlias.equals(ccjsqlTable.getAlisaName())) {
                return ccjsqlTable.getTableName();
            }
        }
        throw new IllegalStateException("Can not find owner from table.");
    }
}
